#!/usr/bin/env python
#  Copyright (c) 2010
#  The Regents of the University of Michigan
#  All Rights Reserved

#  Permission is granted to use, copy, create derivative works, and
#  redistribute this software and such derivative works for any purpose,
#  so long as the name of the University of Michigan is not used in
#  any advertising or publicity pertaining to the use or distribution
#  of this software without specific, written prior authorization. If
#  the above copyright notice or any other identification of the
#  University of Michigan is included in any copy of any portion of
#  this software, then the disclaimer below must also be included.

#  This software is provided as is, without representation or warranty
#  of any kind either express or implied, including without limitation
#  the implied warranties of merchantability, fitness for a particular
#  purpose, or noninfringement.  The Regents of the University of
#  Michigan shall not be liable for any damages, including special,
#  indirect, incidental, or consequential damages, with respect to any
#  claim arising out of or in connection with the use of the software,
#  even if it has been or is hereafter advised of the possibility of
#  such damages.

from pydht import dhtutil, dht
import providence

import Crypto.Cipher.DES3
import event

import bisect
import event
import hashlib
from logging import debug, error, info
import signal
import sys


instances = 0
peers_left = 0
def print_stats():
  print 'Instances: %d' % instances
  print 'Peers left: %d' % peers_left

event.signal(signal.SIGUSR1, print_stats)

def checked_abort():
  if not instances:
    event.abort()

class SingleTorrentManager(object):
  def __init__(self, hash, peers_to_try, download_manager):
    if peers_to_try:
      info('started transfer for %s' % (str(hash).encode('hex'),))
    else:
      info('no peers for %s' % (str(hash).encode('hex')))
      return
    for addr in peers_to_try[:3]:
      download_manager.start_read_transfer(addr, hash, self)

    self.hash = hash
    self.peers_to_try = peers_to_try[3:]
    self.download_manager = download_manager
    self.done = False

    global instances, peers_left
    instances += 1
    peers_left += len(self.peers_to_try)

  def _torrent_finished(self):
    global instances
    global peers_left
    if not self.done:
      self.done = True
      peers_left -= len(self.peers_to_try)
      instances -= 1
      if not instances:
        event.timeout(10, checked_abort)


  def on_error(self):
    if self.done:
      return
    peers = self.peers_to_try
    if not peers:
      self._torrent_finished()
      return
    addr = peers.pop()
    self.download_manager.start_read_transfer(addr, self.hash, self)
    global peers_left
    peers_left -= 1


  def on_finish(self):
    self._torrent_finished()

def packets_overlap(lhs, rhs):
  return (lhs.start < (rhs.start + rhs.length) and
          rhs.start < (lhs.start + lhs.length))


def packet_is_less_than(lhs, rhs):
  # A overlaps B if A's start < B's end and B's start < A's end
  # A < B if A's end < B's start
  # B < A if B's end < A's start
  return (lhs.start + lhs.length) < rhs.start


class Packet(object):
  '''Just the important parts of a data packet.'''
  def __init__(self, data_packet):
    self.start = start = data_packet.start
    self.end = start + data_packet.length
    self.data = data_packet.data

  def __lt__(self, other):
    return self.end <= other.start

  def __eq__(self, other):
    return self.start < other.end and other.start < self.end

  def __repr__(self):
    return 'Packet(%d, %d, %s)' % (self.start, self.end, self.data)

  def __str__(self):
    return 'Packet(%d, %d)' % (self.start, self.end)


MAX_RETRIES = 1
class TransferState(object):
  __slots__ = ['packets', 'total_length', 'torrent_hash', 'torrent_manager', 'peer_address', 'ntries']
  def __init__(self, torrent_hash, peer_address, torrent_manager):
    self.packets = []
    self.total_length = None
    self.torrent_hash = torrent_hash
    self.peer_address = peer_address
    self.torrent_manager = torrent_manager
    self.ntries = 0

  def add_packet(self, packet):
    packets = self.packets
    internal_rep = Packet(packet)
    position = bisect.bisect_left(packets, internal_rep)

    if position < len(packets) and packets[position] == internal_rep:
      return False
    if position and packets[position - 1] == internal_rep:
      return False

    if self.total_length is None:
      self.total_length = packet.total_length

    packets.insert(position, internal_rep)
    return True

  def get_file(self):
    '''Returns the transferred file if the download is finished, else None.'''
    total_length = self.total_length
    if total_length is None:
      return None

    # Check if we have contiguous pieces.
    last_pos = 0
    for packet in self.packets:
      if packet.start != last_pos:
        return None
      last_pos = packet.end
      if last_pos == total_length:
        return ''.join(p.data for p in self.packets)

    return None

  def get_missing_piece_list(self):
    total_length = self.total_length
    if total_length is None:
      return [(0,0)]  # It's all missing.

    # XXX: Shamelessly ripped off from Vuze's DHTTransportUDPImpl.java.
    pos = 0
    pieces = []
    for packet in self.packets:
      if packet.start != pos:
        pieces.append((pos, packet.start - pos))
      pos = packet.end
    if pos != total_length:
      pieces.append((pos, total_length))
    return pieces


class VuzeTorrentDownloader(dht.DHTClient):

  TRANSFER_KEY = hashlib.sha1('org.gudy.azureus2.pluginsimpl.local.ddb.DDBaseTTTorrent').digest()

  def __init__(self, address, finished_callback):
    dht.DHTClient.__init__(self, address)
    self.transfers = {}  # Map connection ID -> transfer state.
    self.finished_callback = finished_callback
    self.went_long_enough = False
    event.timeout(5, self._set_long_enough)

  def _set_long_enough(self):
    self.went_long_enough = True

  def _timeout_transfer(self, conn_id):
    xfers = self.transfers
    state = xfers.get(conn_id)
    if state is not None:
      if state.packets:
        log_str = 'Timeout with %d packets, ' % len(state.packets)
      else:
        log_str = 'Zero-packet timeout, '

      if state.ntries == MAX_RETRIES:
        state.torrent_manager.on_error()
        del xfers[conn_id]
        log_str += 'gave up'
      else:
        pieces = state.get_missing_piece_list()
        dhthash = hashlib.sha1(state.torrent_hash).digest()
        for pos, length in pieces:
          self.send_data_request(state.peer_address,
                                 VuzeTorrentDownloader.TRANSFER_KEY,
                                 dhthash,
                                 conn_id,
                                 pos,
                                 length)
        state.ntries += 1
        log_str += 'retrying (try %d)' % state.ntries
      debug(log_str)

  def start_read_transfer(self, other_address, torrent_hash, torrent_manager):
    conn_id = self._get_connection_id()
    while conn_id in self.transfers:  # Just in case.
      conn_id = self._get_connection_id()

    self.transfers[conn_id] = TransferState(torrent_hash, other_address,
                                            torrent_manager)

    self.send_data_request(other_address, VuzeTorrentDownloader.TRANSFER_KEY,
                           hashlib.sha1(torrent_hash).digest(),
                           conn_id)

  def handle_data(self, packet, sender, *args):
    conn_id = packet.header.connection_id
    transfer_state = self.transfers.get(conn_id)
    if transfer_state is None:
      debug('Got irrelevant packet from %s with conn_id %s' % (str(sender), hex(conn_id)))
      debug(repr(packet))
      return

    debug('Got packet from %s' % str(sender))
    if transfer_state.add_packet(packet):
      file = transfer_state.get_file()
      if file is not None:
        info('Download of %s terminated' % str(transfer_state.torrent_hash).encode('hex'))
        del self.transfers[conn_id]
        if file[0]:
          file = decrypt_torrent(file, transfer_state.torrent_hash)
          info('Successful download of %s' %
               str(transfer_state.torrent_hash).encode('hex'))
          self.finished_callback(transfer_state.torrent_hash, file)
          transfer_state.torrent_manager.on_finish()
        else:
          error('Download of %s corrupted!' % str(transfer_state.torrent_hash).encode('hex'))
          transfer_state.torrent_manager.on_error()

        if not self.transfers and self.went_long_enough:
          event.abort()
    else:
      debug('discarded packet as duplicate')

  def _actually_send(self, packet, address, conn_id, *args):
    dht.DHTClient._actually_send(self, packet, address, conn_id, *args)
    event.timeout(60, self._timeout_transfer, conn_id)

def decrypt_torrent(data, torrent_hash):
  should_decrypt = ord(data[1])
  data = data[2:]
  if should_decrypt:
    remainder = len(data) % 8
    data += '\x00' * ((8 - remainder) % 8)
    cipher = Crypto.Cipher.DES3.new(torrent_hash + '\x00' * 4,
                                    Crypto.Cipher.DES3.MODE_ECB)
    return cipher.decrypt(data)
  else:
    return data

def main(argv):
  dhtutil.HACK_DISABLE_TIMER = False
  def finished(torrent_hash, file):
    print 'Got file for hash %s' % str(torrent_hash).encode('hex')
    if ord(file[0]) != 1:
      print 'Bad crypto version!'
    else:
      if ord(file[1]):
        file = decrypt_torrent(file[2:], torrent_hash)
      else:
        file = file[2:]
      fetched = open('%s.out' % str(torrent_hash).encode('hex'), 'w')
      print >> fetched, file
      event.abort()

  downloader = VuzeTorrentDownloader((argv[1], int(argv[2])), finished)
  downloader.start_read_transfer((argv[3], int(argv[4])), argv[5].decode('hex'))

  event.signal(signal.SIGTERM, event.abort)
  event.signal(signal.SIGINT, event.abort)
  event.dispatch()



if __name__ == '__main__':
  main(sys.argv)
